"""

From: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix

Test the CycleGAN on the target dataset

"""

import os

from options.test_options import TestOptions
from data import create_dataset
from models import create_model
from pathlib import Path
from tqdm import tqdm
from util.util import tensor2im
from detectron2.utils.visualizer import Visualizer as vis

if __name__ == '__main__':
    opt = TestOptions().parse()     # get test options
    opt.num_threads = 0             # test code only supports num_threads = 0
    opt.batch_size = 1              # test code only supports batch_size = 1
    opt.serial_batches = True       # disable data shuffling; comment this line if results on randomly chosen images are needed.
    opt.no_flip = True              # no flip; comment this line if results on flipped images are needed.
    opt.display_id = -1             # no visdom display; the test code saves the results to a HTML file.

    dataset = create_dataset(opt)   # create a dataset given opt.dataset_mode and other options
    model = create_model(opt)       # create a model given opt.model and other options
    model.setup(opt)                # regular setup: load and print networks; create schedulers

    target_path = os.path.join(opt.results_dir, opt.name)
    target_A_path = os.path.join(target_path, 'A')
    target_B_path = os.path.join(target_path, 'B')    
    Path(target_path).mkdir(parents=True, exist_ok=True)
    Path(target_A_path).mkdir(parents=True, exist_ok=True)
    Path(target_B_path).mkdir(parents=True, exist_ok=True)

    # test with eval mode. This only affects layers like batchnorm and dropout.
    if opt.eval:
        model.eval()

    # for each data
    for i, data in enumerate(tqdm(dataset)):

        model.set_input(data)  # unpack data from data loader
        model.test()           # run inference
        visuals = model.get_current_visuals()  # get image results

        img_name_A = os.path.basename(data['A_paths'][0]).replace(".jpg", "")
        img_name_B = os.path.basename(data['B_paths'][0]).replace(".jpg", "")

        for label, image in visuals.items():
            if 'fake' in label:
                image_numpy = tensor2im(image)
                v = vis(image_numpy)
                out = v.get_output()

                if 'B' in label:
                    out.save(os.path.join(target_A_path, "%s.jpg" % img_name_A))
                elif 'A' in label:
                    out.save(os.path.join(target_B_path, "%s.jpg" % img_name_B))